# load packages
library(tidyverse)
library(brms)
library(tidybayes)
library(hrbrthemes)
library(kableExtra)

fit202 <- read_rds("output/fits/model-01-structure-2-0-2.rds") 

summary(fit202)

# a formal test of the interaction hypothesis

test <- spread_draws(fit202, `b_treat_indicator_std:awareness_std`) %>%
  select(prod = `b_treat_indicator_std:awareness_std`) %>%
  summarize(posterior_probability = mean(prod < 0),
            posterior_mean = mean(prod)) %>%
  glimpse()

x_tmp <- paste0("A one standard deviation increase in awareness decreases the treatment effect by about ",
                -round(test$posterior_mean, 2),  
                " units, a two standard deviation increase by about ", 
                -round(2*test$posterior_mean, 2), 
                " units, and a minimum-to-maximum increase by about ", 
                "[see below]", 
                " units (from a treatment effect of ", 
                "[see below]", 
                " to ", 
                "[see below]", 
                "). The posterior probability that the coefficient for the product term is negative is ", 
                round(test$posterior_probability, 2), 
                "--moderate evidence for our interaction hypothesis.\n\n"); x_tmp
cat(x_tmp,
    file = "output/qis.txt")

df <- read_rds("output/rescaled-data.rds") %>%
  glimpse()

# create fake data set
fake_df <- crossing(treat_indicator = c(0, 1),
                    policy = unique(df$policy)) %>%
  left_join(distinct(select(df, policy, stem, awareness, awareness_std, issue, category))) %>%
  mutate(treat_indicator_std = treat_indicator - 0.5) %>%
  glimpse()

# simulate quantities of interest
qi_df <- fake_df %>%
  add_epred_draws(fit202) %>%
  ungroup() %>% 
  select(-ends_with("_std")) %>%
  group_by(policy, stem, issue, category, awareness, .draw) %>%
  select(-.chain, -.iteration, -.row) %>%
  spread(treat_indicator, .epred) %>%
  mutate(treatment_effect = `1` - `0`) %>%
  select(-`0`, -`1`) %>% 
  select(policy, stem, issue, category, awareness, draw = .draw, treatment_effect) %>%
  ungroup() %>%
  mutate(stem = reorder(stem, treatment_effect)) %>%
  mutate(category = fct_recode(category, 
                               `Economic Policy` = "Economic",
                               `Social Policy` = "Social")) %>%
  glimpse()

# summaries qis
qi_sum_df <- qi_df %>%
  group_by(policy, stem, issue, category, awareness) %>%
  summarize(post_avg = mean(treatment_effect), 
            post95 = quantile(treatment_effect, 0.95), 
            post05 = quantile(treatment_effect, 0.05), 
            prob_gt0 = mean(treatment_effect > 0),
            ht_result = case_when(prob_gt0 > 0.95 ~ "Strong Evidence", 
                                  prob_gt0 > 0.90 ~ "Moderate Evidence", 
                                  TRUE ~ "Weak or No Evidence"),
            ht_result = reorder(ht_result, prob_gt0)) %>%
  ungroup() %>%
  mutate(stem = reorder(stem, post_avg)) %>%
  glimpse()

# plot treatment effect estimates and cis
gg <- ggplot(qi_sum_df, aes(x = post_avg, y = stem, 
                            xmin = post05, xmax = post95,
                            color = awareness)) + 
  geom_errorbarh(height = 0) + 
  geom_point() + 
  theme_bw() + theme(legend.position = "bottom", legend.title.align = 1) + 
  #theme_ipsum() + theme(legend.position = "bottom", legend.title.align = 1) + 
  #facet_grid(rows = vars(category), scales = "free_y", space = "free") + 
  labs(title = "Estimates of Treatment Effect for Each Policy",
       subtitle = "Posterior Average and 90% Percentile Interval",
       y = NULL, #"Policy Stems", 
       x = "Treatment Effect",
       color = "Percent Aware of Parties' Position") + 
  facet_grid(rows = vars(category), scales = "free_y", space = "free") + 
  scale_color_gradient(low = "#d95f02", high = "#1b9e77", labels = scales::percent_format(accuracy = 1)); gg
ggsave("figs/fig1-policy-effects.tiff",
       width = 12, height = 12, scale = 1.0)

# plot treatment effect estimates and cis
ggplot(qi_sum_df, aes(y = post_avg, x = awareness, 
                            ymin = post05, ymax = post95,
                            color = category, shape = category)) + 
  geom_hline(yintercept = 0) + 
  geom_errorbar() + 
  geom_point(fill = "white") + 
  scale_color_manual(values = c("#1b9e77", "#d95f02", "#7570b3")) +
  scale_x_continuous(labels = scales::percent_format(accuracy = 1)) + 
  scale_y_continuous(breaks = seq(0, 1, by = 0.2)) + 
  scale_shape_manual(values = c(16, 21, 4)) + 
  #geom_smooth(se = FALSE, color = "grey90", aes(color = NULL, shape = NULL)) + 
  #theme_ipsum() + theme(legend.position = "bottom", legend.title.align = 1) + 
  theme_bw() + theme(legend.position = "bottom", legend.title.align = 1) + 
  labs(title = "Estimates of Treatment Effect for Each Policy as Awareness Varies",
       subtitle = "Posterior Average and 90% Percentile Interval",
       x = "Percent Aware of Parties' Positions on the Policy", 
       y = "Treatment Effect", 
       color = "Category",
       shape = "Category") 
ggsave("figs/fig2-policy-effect-awareness-scatter.tiff",
       width = 8, height = 6, scale = 1.0)

# correlation between treatment effects and awareness
r_tmp <- cor(qi_sum_df$post_avg, qi_sum_df$awareness)
cat(paste0("The scatterplot clearly shows a negative relationship between treatment effect size and the level of awareness: ", round(r_tmp, 2), "\n\n"),
           file = "output/qis.txt", append = TRUE)

# treatment effects table
et <- qi_sum_df %>%
  select(category, issue, stem, awareness,post_avg, post95, post05, prob_gt0) %>%
  mutate(stem = reorder(stem, awareness)) %>%
  mutate(issue = reorder(issue, awareness)) %>%
  mutate(category = reorder(category, awareness)) %>%
  arrange(category, issue, stem) %>%
  mutate(Category = category,
         Issue = issue,
         Stem = stem, 
         `Percent Aware of Parties' Positions` = scales::percent(awareness, accuracy = 1), 
         `Posterior Average` = round(post_avg, 2), 
         `90% Credible Interval` = paste0("[", round(post05, 2), ", ", round(post95, 2), "]"),
         `Percent of Posterior Above Zero` = scales::percent(prob_gt0, accuracy = 1)) %>%
  glimpse() %>%
  select(Category:`Percent of Posterior Above Zero`) %>% 
  write_rds("figs/taba2-effects-table.rds")